#include "leafsegmentation.h"
#include "leafcurvature.h"
#include "leafdatabase.h"

#include <opencv2/core/core.hpp>
#include <opencv2/highgui/highgui.hpp>
#include <opencv2/imgproc/imgproc.hpp>

#include <iostream>
#include <fstream>

using namespace cv;
using namespace std;

Mat getSegmentedImg(const Mat& leaf_img, unsigned int num_segments) {
    imwrite( "original.jpg", leaf_img );

    // visualize boundary
    Mat white_img(leaf_img.size(), CV_8UC3);white_img.setTo(Vec3b(255,255,255));
    Mat black_img(leaf_img.size(), CV_8UC3);black_img.setTo(Vec3b(0,0,0));

    std::vector<const Mat*> white_on_black_images;
    white_on_black_images.push_back(&black_img);
    white_on_black_images.push_back(&white_img);

    std::vector<const Mat*> img_on_black_images;
    img_on_black_images.push_back(&black_img);
    img_on_black_images.push_back(&leaf_img);

    ImageMask leaf_mask=LeafSegmentation::segment(leaf_img, 2, num_segments);

    Mat segmented = leaf_mask.applyTo(&white_img);
    namedWindow("segmented", CV_WINDOW_AUTOSIZE);
    imshow("segmented", segmented);

    imwrite( "segmented.jpg", segmented );

    ImageMask cleaned_mask=LeafSegmentation::removeFalsePositive(leaf_mask,0.3,2);

    Mat cleaned = cleaned_mask.applyTo(&white_img);
    namedWindow("cleaned", CV_WINDOW_AUTOSIZE);
    imshow("cleaned", cleaned);

    imwrite( "cleaned.jpg", cleaned );

    Scalar area=sum(cleaned_mask.mMask);
    //if((float)area/(float)(destemed_mask.mMask.size().area)<)
    int num_cluster=2;
    while(*(area.val)<50&&num_cluster<10) {
        num_cluster++;
        leaf_mask=LeafSegmentation::segment(leaf_img, num_cluster, num_segments);

        segmented = leaf_mask.applyTo(&white_img);
        namedWindow("segmented2", CV_WINDOW_AUTOSIZE);
        imshow("segmented2", segmented);

        imwrite( "segmented"+QString::number(num_cluster).toStdString()+".jpg", segmented );

        cleaned_mask=LeafSegmentation::removeFalsePositive(leaf_mask,0.1,2);

        cleaned = cleaned_mask.applyTo(&white_img);
        namedWindow("cleaned2", CV_WINDOW_AUTOSIZE);
        imshow("cleaned2", cleaned);

        imwrite( "cleaned"+QString::number(num_cluster).toStdString()+".jpg", cleaned );

        area=sum(cleaned_mask.mMask);
    }

    ImageMask destemed_mask=LeafSegmentation::removeStem(cleaned_mask);

    Mat destemed = destemed_mask.applyTo(&white_img);
    //imwrite("top_hat.png", destemed);
    namedWindow("destemed", CV_WINDOW_AUTOSIZE);
    imshow("destemed", destemed);

    imwrite( "destemed.jpg", destemed );

    waitKey(0); // Wait for a keystroke in the window
    exit(0);

    ImageMask boundary_mask=LeafCurvature::boundary(leaf_mask.mMask);

    char buffer[100];
    for(int radius=5;radius<=40;radius*=2){
        sprintf(buffer,"%d",radius);
        string radius_str=buffer;

        std::string arclength_histogram_filename="curvature_arclength_"+radius_str+".dat";
        ofstream arclength_histogram_file(arclength_histogram_filename.c_str());
        std::vector<int> arclength_histogram=LeafCurvature::arclengthHistogram(leaf_mask,radius,15);
        for(unsigned int i=0;i<arclength_histogram.size();i++)
            arclength_histogram_file<<arclength_histogram[i]<<std::endl;
        arclength_histogram_file.close();

        std::string area_histogram_filename="curvature_area_"+radius_str+".dat";
        ofstream area_histogram_file(area_histogram_filename.c_str());
        std::vector<int> area_histogram=LeafCurvature::areaHistogram(leaf_mask,radius,15);
        for(unsigned int i=0;i<area_histogram.size();i++)
            area_histogram_file<<area_histogram[i]<<std::endl;
        area_histogram_file.close();
    }
    
    //visualize segmentation

    Mat boundary_img=boundary_mask.applyTo(white_on_black_images);
    imwrite( "boundary.jpg", boundary_img );
    Mat segment_img=leaf_mask.applyTo(white_on_black_images);
    imwrite( "segmentation.jpg", segment_img );
    Mat segment_img2=leaf_mask.applyTo(img_on_black_images);
    imwrite( "segmentation2.jpg", segment_img2 );

    Mat output=boundary_mask.applyTo(white_on_black_images);
    return output;
}

int main( int argc, char** argv )
{   
    string plant_dir;
    string img_name;
    if(argc>1){
        img_name=argv[1];
    }
    else{
//        // test removing stem
//        plant_dir="../../leafsnap-dataset/dataset/images/field/acer_platanoides";
//        img_name = plant_dir + "/1249060544_0000.jpg";


//        // test remove false positive
//        plant_dir="../../leafsnap-dataset/dataset/images/field/abies_concolor";
//        img_name = plant_dir + "/12995309740671.jpg";

//        // test remove false positive
//        plant_dir="../../leafsnap-dataset/dataset/images/field/abies_concolor";
//        img_name = plant_dir + "/12995309740678.jpg";

//        // test remove false positive
//        plant_dir="../../leafsnap-dataset/dataset/images/field/acer_ginnala";
//        img_name = plant_dir + "/13291762516235.jpg";
        


//                // test remove false positive
//                plant_dir="../../leafsnap-dataset/dataset/images/field/quercus_rubra";
//                img_name = plant_dir + "/1249060144_0003.jpg";

//                // test remove false positive
//                plant_dir="../../leafsnap-dataset/dataset/images/lab/aesculus_hippocastamon";
//                img_name = plant_dir + "/ny1017-10-3.jpg";

//        // test remove false positive
//        plant_dir="../../leafsnap-dataset/dataset/images/field/albizia_julibrissin";
//        img_name = plant_dir + "/13001200140350.jpg";

//        //test databse
//        plant_dir="../../leafsnap-dataset/dataset/images/field/abies_concolor";
//        img_name = plant_dir + "/12995307070714.jpg";
//        plant_dir="../../leafsnap-dataset/dataset/images/field/abies_nordmanniana";
//        img_name = plant_dir + "/13291651120806.jpg";

        // test remove false positive
        plant_dir="../../bad_images";
        img_name = plant_dir + "/IMG_0758.JPG";
    }

    Mat leaf_img = imread(img_name,CV_LOAD_IMAGE_COLOR);
    pyrDown( leaf_img, leaf_img, Size( leaf_img.cols/2, leaf_img.rows/2 ));
    pyrDown( leaf_img, leaf_img, Size( leaf_img.cols/2, leaf_img.rows/2 ));
//    ImageMask leaf_mask=LeafSegmentation::segment(leaf_img, 2);

//    std::vector<float> arclength_feature;
//    std::vector<float> area_feature;
//    for(int radius=2;radius<=50;radius+=2){
//        std::vector<float> arclength_histogram=LeafCurvature::normalizedArclengthHistogram(leaf_mask.mMask,radius,21);
//        arclength_feature.insert(arclength_feature.end(),arclength_histogram.begin(),arclength_histogram.end());
//        std::vector<float> area_histogram=LeafCurvature::normalizedAreaHistogram(leaf_mask.mMask,radius,21);
//        area_feature.insert(area_feature.end(),area_histogram.begin(),area_histogram.end());
//    }
//    LeafDatabase& database=*LeafDatabase::GetInstance();

//    database.closestNEntries(arclength_feature,area_feature);

    Mat output = getSegmentedImg(leaf_img, 2);
    
    namedWindow("test", CV_WINDOW_AUTOSIZE);
    imshow("test", output);

    waitKey(0); // Wait for a keystroke in the window
    return 0;
}



